import torch
import torchvision
import torch.nn as nn
import torch.utils.data as Data
from model import LeNet
# 导入设定好的模型Lenet
model = LeNet()
# 设定遍历轮次，批次大小，学习率lr
Epoch = 5
batch_size = 64
lr = 0.001
# 导入训练集
train_data = torchvision.datasets.MNIST(root='./data/', train=True, transform=torchvision.transforms.ToTensor(), download=False)
# 将数据喂给模型
train_loader = Data.DataLoader(train_data, batch_size=batch_size, shuffle=True, num_workers=0, drop_last=True)
# 定义损失函数为交叉熵
loss_function = nn.CrossEntropyLoss()
# 优化器使用Adam
optimizer = torch.optim.Adam(model.parameters(), lr=lr)
# 启用pytorch的自动求导机制，并利用其实现反向传播和参数更新
torch.set_grad_enabled(True)
# 启用Batch Normalization层和Dropout层
model.train()
# 启用CUDA加速
# torch.cuda.is_available()可以判断本设备是否支持CUDA，如果支持就返回True，不支持就返回False。
# 有了这个函数，我们就可以让其自动判断是否支持CUDA加速并自动选择设备了。而不支持CUDA加速的设备我们可以使用cpu来进行。
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
model.to(device)
# 模型训练
# 在pytorch中，神经网络的训练一般是分以下几个步骤进行的：
# 1) 获得DataLoader中的数据x和标签y
# 2) 将优化器的梯度清零
# 3) 将数据送入模型中获得预测的结果y_pred
# 4) 将标签和预测结果送入损失函数获得损失
# 5) 将损失值反向传播
# 6) 使用优化器对模型的参数进行更新

for epoch in range(Epoch):
    # 定义损失值和准确率
    running_loss = 0.0
    acc = 0.0
    # 每一次迭代实现
    for step, data in enumerate(train_loader):
        # 将data中的数据和标签取出，其中数据为x，标签为y
        x, y = data
        # 调用optimizer.zero_grad()将优化器的梯度清零
        optimizer.zero_grad()
        # 将数据送入模型获得预测结果y_pred
        y_pred = model(x.to(device, torch.float))
        # 将标签和预测结果送入损失函数获得损失，类型变化为一致
        loss = loss_function(y_pred, y.to(device, torch.long))
        # 将损失值反向传播，这边直接调用loss.backward()即可
        loss.backward()
        # 累加值，计算迭代100次的平均值
        running_loss += float(loss.data.cpu())
        # 获取相关性最大的下标
        pred = y_pred.argmax(dim=1)
        # 计算预测正确的数量
        acc += (pred.data.cpu() == y.data).sum()
        # 优化器对模型的参数进行更新
        optimizer.step()
        # 判断迭代次数是否满足输出平均值
        if step % 100 == 99:
            loss_avg = running_loss / (step + 1)
            # 预测正确的数量，将其除以当前已预测图像的总数即可，总数量可以通过step * batch_size求得
            acc_avg = float(acc / ((step + 1) * batch_size))
            print('Epoch', epoch + 1, ',step', step + 1, '| Loss_avg: %.4f' % loss_avg, '|Acc_avg:%.4f' % acc_avg)
# 模型保存
torch.save(model, './LeNet.pkl')